#coding=utf-8

from cigeratte_utils import *

class GenSample:
    def __init__(self, font_ch = font_ch_plate):
        self.my_font =  ImageFont.truetype(font_ch, 49, 0)
        self.bg  = cv2.imread(plate_template_img_path)
        self.smu = cv2.imread(plate_noise_img_path)

    def generate(self,text):
        if len(text) == cigeratte_chars_count:
            complex_img = np.zeros((plate_height_default, plate_width_default, 3), dtype=np.uint8)
            text_img = draw_text_3(self.my_font,text)
            rimg = rotate_img(text_img,rand_int(6,-6),1)
            # rimg = Erosion(rimg,1)
            size = [rimg.shape[0],rimg.shape[1]]

            pos = ( (complex_img.shape[0] - size[0])/2,(complex_img.shape[1] - size[1])/2 )
            pos = (max(pos[0],0),max(pos[1],0))
            start = [0,0]
            size_r = size
            if complex_img.shape[0] < size[0]:
                start[0] = (size[0] - complex_img.shape[0]) / 2
                size_r[0] = complex_img.shape[0]
            complex_img[ pos[0]: size[0]+pos[0] , pos[1]:(size[1]+pos[1]) ] = rimg[start[0]:size_r[0]+start[0],:]

            complex_img = cv2.bitwise_or(complex_img, self.bg)

            complex_img = append_noise(complex_img,self.smu)
            # self.img = cv2.cvtColor( cv2.bitwise_or(self.img,self.bg) , cv2.COLOR_BGR2RGB)
            return complex_img

    @staticmethod
    def gen_rand_string():
        plateStr = ''
        for cpos in range(cigeratte_chars_count):
            if cpos < 16 or cpos >= 20:
                plateStr += PLATE_CHARS_TOTAL[rand_int(10)]
            else:
                plateStr += PLATE_CHARS_TOTAL[rand_int(36,10)]

        return plateStr

    def gen_one(self,size = (plate_width_default,plate_height_default), channels = 1):
        text = self.gen_rand_string()
        # text = text.replace(text[:5], "70928")
        # text = text.replace(text[16:22],"GZYC44")
        img = self.generate(text)
        # img = cv2.resize(img, dsize = (size[0]/2,size[1]/2), interpolation=cv2.INTER_AREA)

        if channels == 1:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        return text, img

    def gen_batch(self, batchSize, outputPath, size):
        if not os.path.exists(outputPath):
            os.makedirs(outputPath)
        for i in xrange(batchSize):
            text , img = self.gen_one(size)
            filename = os.path.join(outputPath, str(i).zfill(5) + '_' + text + ".jpg")
            cv2.imwrite(filename, img)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--font_ch', default = font_ch_plate)
    # parser.add_argument('--bg_dir', default = plate_bg_dir )
    parser.add_argument('--out_dir', default = 'cigeratte_tf/train', help = 'output dir')
    parser.add_argument('--make_num', default = 20, type = int, help = 'num')
    parser.add_argument('--img_w', default= plate_width_default, type = int, help='num')
    parser.add_argument('--img_h', default= plate_height_default, type = int, help='num')
    return parser.parse_args()

def main(args):
    g = GenSample(args.font_ch)
    g.gen_batch(args.make_num, args.out_dir, (args.img_w, args.img_h))

if __name__ == '__main__':
    main(parse_args())
